#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
会话管理API路由
"""

import logging
import uuid
import json
from datetime import datetime
from bson import ObjectId
from flask import Blueprint, request, jsonify
from db.mongo import get_db
from base.services.processor import extract_services_from_trace, extract_requests_from_trace

# 自定义JSON编码器，处理MongoDB的ObjectId
class JSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, ObjectId):
            return str(obj)
        if isinstance(obj, datetime):
            return obj.isoformat()
        return json.JSONEncoder.default(self, obj)

logger = logging.getLogger(__name__)
sessions_bp = Blueprint('sessions', __name__)

@sessions_bp.route('', methods=['POST'])
def create_session():
    """创建新的测试会话"""
    data = request.json

    if not data or 'trace_data' not in data:
        return jsonify({'error': 'Missing trace_data'}), 400

    try:
        # 生成唯一会话ID
        session_id = str(uuid.uuid4())

        # 创建会话记录
        session = {
            'session_id': session_id,
            'created_at': datetime.now(),
            'trace_data': data['trace_data'],
            'top_level_request': data.get('top_level_request', {})
        }

        # 保存会话
        db = get_db()
        db.sessions.insert_one(session)

        # 从trace数据中提取服务信息
        services = extract_services_from_trace(data['trace_data'])

        # 保存服务信息
        for service in services:
            service['session_id'] = session_id
            db.services.insert_one(service)

        # 从trace数据中提取请求信息
        requests_info = extract_requests_from_trace(data['trace_data'], services)

        # 保存请求信息
        for req in requests_info:
            req['session_id'] = session_id
            db.requests.insert_one(req)

        logger.info(f"创建会话成功: {session_id}")

        # 使用自定义JSON编码器处理MongoDB对象
        result = {
            'session_id': session_id,
            'services': services,
            'requests': requests_info
        }
        return json.dumps(result, cls=JSONEncoder), 201, {'Content-Type': 'application/json'}

    except Exception as e:
        logger.error(f"创建会话失败: {e}")
        return jsonify({'error': str(e)}), 500

@sessions_bp.route('', methods=['GET'])
def list_sessions():
    """获取所有会话列表"""
    try:
        db = get_db()
        sessions = list(db.sessions.find({}, {
            '_id': 0,
            'session_id': 1,
            'created_at': 1
        }))

        return json.dumps(sessions, cls=JSONEncoder), 200, {'Content-Type': 'application/json'}
    except Exception as e:
        logger.error(f"获取会话列表失败: {e}")
        return jsonify({'error': str(e)}), 500

@sessions_bp.route('/<session_id>', methods=['GET'])
def get_session(session_id):
    """获取特定会话的详细信息"""
    try:
        db = get_db()
        session = db.sessions.find_one({'session_id': session_id}, {'_id': 0})

        if not session:
            return jsonify({'error': 'Session not found'}), 404

        # 获取会话相关的服务信息
        services = list(db.services.find({'session_id': session_id}, {'_id': 0}))
        session['services'] = services

        return json.dumps(session, cls=JSONEncoder), 200, {'Content-Type': 'application/json'}
    except Exception as e:
        logger.error(f"获取会话详情失败: {e}")
        return jsonify({'error': str(e)}), 500

@sessions_bp.route('/<session_id>', methods=['DELETE'])
def delete_session(session_id):
    """删除特定会话及其相关数据"""
    try:
        db = get_db()

        # 检查会话是否存在
        session = db.sessions.find_one({'session_id': session_id})
        if not session:
            return jsonify({'error': 'Session not found'}), 404

        # 删除会话及相关数据
        db.sessions.delete_one({'session_id': session_id})
        db.services.delete_many({'session_id': session_id})
        db.requests.delete_many({'session_id': session_id})
        db.responses.delete_many({'session_id': session_id})

        logger.info(f"删除会话成功: {session_id}")
        return jsonify({'message': 'Session deleted successfully'}), 200
    except Exception as e:
        logger.error(f"删除会话失败: {e}")
        return jsonify({'error': str(e)}), 500
